import logging
import subprocess

import pexpect

from tools.log_custom import setup_logging

setup_logging(logging.INFO)


class HostCmdTool:
    def __init__(self):
        pass

    @staticmethod
    def copy_ssh_id(
            target_ip, target_ssh_port="22", target_user="root", target_pwd=""
    ):
        """
        ssh-copy-id -i ~/.ssh/id_ed25519 -p 22 root@ip
        :return:
        """
        cmd = f"ssh-copy-id -p {target_ssh_port} {target_user}@{target_ip}"
        logging.info(cmd)
        child = pexpect.spawn(cmd, encoding="utf-8")
        index = child.expect(
            ["password", "exist", pexpect.EOF, pexpect.TIMEOUT], timeout=10
        )
        if index == 0:
            child.sendline(target_pwd)
            i = child.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=10)
            if i == 0:
                cmd_show_data = child.before
                cmd_output = cmd_show_data.split("\r\n")
                for data in cmd_output:
                    if "Number of key" in data:
                        logging.info(data)
            elif i == 1:
                logging.error("TIMEOUT")
                child.kill(0)
        elif index == 1:
            logging.warning(
                "All keys were skipped because they already exist on the remote system."
            )
        elif index == 2:
            logging.error("EOF")
            child.kill(0)
        elif index == 3:
            logging.error("TIMEOUT")
            child.kill(0)

    @staticmethod
    def rsync_upload(
            target_ip,
            target_user="root",
            local_dir="~/upload/*",
            upload_dir="/root/test/",
    ):
        """
        在本机上执行rsync命令，上传文件
        rsync -avP ~/test/* root@ip:/root/test/
        :param target_ip: ip
        :param target_user: root
        :param local_dir: ~/upload/*
        :param upload_dir: /root/test/
        """
        cmd = f"rsync -avP {local_dir} {target_user}@{target_ip}:{upload_dir}"
        logging.info(cmd)
        child = pexpect.spawn(cmd, encoding="utf-8")
        index = child.expect(
            ["total size", "No such file or directory", pexpect.EOF, pexpect.TIMEOUT],
            timeout=10,
        )
        cmd_show_data = child.before
        cmd_output = cmd_show_data.split("\r\n")
        if index == 0:
            for data in cmd_output:
                logging.info(data)
        elif index == 1:
            logging.error("No such file or directory")
            for data in cmd_output:
                logging.error(data)
        elif index == 2:
            logging.error("EOF")
            child.kill(0)
        elif index == 3:
            logging.error("TIMEOUT")
            child.kill(0)

    @staticmethod
    def rsync_download(
            server_ip,
            server_user="root",
            server_dir="/root/test/*",
            local_dir="~/download/",
    ):
        """
        在本机上执行rsync命令，下载文件
        rsync -avP root@ip:/root/test/* ~/test/
        :param server_ip: IP
        :param server_user: 默认root
        :param server_dir: /root/test/*
        :param local_dir: ~/download/
        :return:
        """
        cmd = f"rsync -avP {server_user}@{server_ip}:{server_dir} {local_dir}"
        logging.info(cmd)
        child = pexpect.spawn(cmd, encoding="utf-8")
        index = child.expect(
            ["total size", "No such file or directory", pexpect.EOF, pexpect.TIMEOUT],
            timeout=10,
        )
        cmd_show_data = child.before
        cmd_output = cmd_show_data.split("\r\n")
        if index == 0:
            for data in cmd_output:
                logging.info(data)
        elif index == 1:
            logging.error("No such file or directory")
            for data in cmd_output:
                logging.error(data)
        elif index == 2:
            logging.error("EOF")
            child.kill(0)
        elif index == 3:
            logging.error("TIMEOUT")
            child.kill(0)

    @staticmethod
    def run_script(script_path="./test.sh"):
        """
        在本机上执行bash命令
        :param filename: 要执行的命令字符串
        """
        logging.info("----------------------- 执行脚本 ----------------------")
        logging.info(script_path)

        process = subprocess.Popen(
            ["bash", script_path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            shell=True,
        )

        # 实时读取输出
        if process.stdout:
            logging.info("----------------------- 输出 ----------------------")
            for line in process.stdout:
                logging.info(line.strip())

        if process.stderr:
            logging.error("----------------------- 错误信息 ----------------------")
            for line in process.stderr:
                logging.error(line.strip())

        # 等待进程结束
        process.wait()

    @staticmethod
    def run_sh(command):
        """
        执行命令并返回结果
        """
        logging.info("----------------------- 命令 ----------------------")
        logging.info(command)
        result = subprocess.run(
            command,
            shell=True,  # 启用管道
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True
        )
        if result.stdout:
            logging.info("----------------------- 输出 ----------------------")
            logging.info(result.stdout)
        if result.stderr:
            logging.error("----------------------- 错误信息 ----------------------")
            logging.error(result.stderr)


if __name__ == "__main__":
    host = HostCmdTool()
    host.run_sh("pwd")
